from pathlib import Path
from utils.runer import gen_green_kwargs, main
from lightning.pytorch import seed_everything
from utils.l2_loss import l2_loss
import torch.nn.functional as F
import json

hyper_kwargs = {
    'net_kwargs' : {
        "VaryUNet": {
            "features": [16],
            "layers": [6],
        }
    },
    'gridSize' : [256],
    'batch_size' : [5],
    'method' : ['imconv'],
    'maxiter' : [5],
    'lr' : [1e-3], 
    'max_epochs': 150,
    'gpus' : 1,
    'data_path' : './data',
    'loss_fn' : [F.mse_loss],
    'data_type' : ['xyf'],
    'equation_type' : ['Lap'],
    'TrainN':[2000],
    'js_log_path': './js_logs/data_amount.json',
    'K_strategy' : ['const']}

for kwargs, log_info in gen_green_kwargs(**hyper_kwargs):
    js_log_path = Path(log_info['js_log_path'])
    if not js_log_path.exists():
        js = {}
    else:
        with open(js_log_path, 'r') as f:
            js = json.load(f,)

    seed = seed_everything(0)
    if not log_info['exp_name'] in js.keys():
        js[log_info['exp_name']] = {
            'epochs':0,
            'fit_start_time':0,
            'fit_end_time':0,
            'fit_spend_time':0,
            'best_model_ckpt':'_',
            'last_model_ckpt':'_',
            'exception':None,
            'seed':seed
        }
        with open(js_log_path, 'w') as f:
            log_data = json.dumps(js)
            f.write(log_data)
    else:
        if js[log_info['exp_name']]['exception'] is None:
            if js[log_info['exp_name']]['epochs'] == hyper_kwargs['max_epochs']-1:
                continue
            else:
                last_ckpt_path = Path(js[log_info['exp_name']]["last_model_ckpt"])
                if last_ckpt_path.is_file():
                    kwargs.update({'ckpt': last_ckpt_path})
                else:
                    best_ckpt_path = Path(js[log_info['exp_name']]["best_model_ckpt"])
                    if best_ckpt_path.is_file():
                        kwargs.update({'ckpt': best_ckpt_path})
        else:
            js[log_info['exp_name']]['exception'] = None
            with open(js_log_path, 'w') as f:
                log_data = json.dumps(js)
                f.write(log_data)
    try:
        print('\n', log_info['exp_name'], '\n')
        main(**kwargs)
        
    
    except Exception as e:
        print(e)
        with open(js_log_path, 'r') as f:
            js = json.load(f)
        js[log_info['exp_name']]['exception'] = str(e)
        with open(js_log_path, 'w') as f:
            f.write(json.dumps(js))
        continue
    